package com.agentframework.core.provider;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.image.ImageModel;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
 * @author wangjun
 **/
public class DefaultProvider implements Provider, ApplicationListener<ApplicationReadyEvent> {

    private final Map<String, ChatModel> chatModelMap = Collections.synchronizedMap(new HashMap<>());

    private final Map<String, ImageModel> imageModelMap = Collections.synchronizedMap(new HashMap<>());

    private final Map<String, EmbeddingModel> embeddingModelMap = Collections.synchronizedMap(new HashMap<>());

    @Override
    public ChatModel getChatModel() {
        return getChatModel(LlmType.DEEPSEEK);
    }

    @Override
    public ChatModel getChatModel(LlmType llmType) {
        return chatModelMap.get(llmType.getType() == null ? LlmType.DEEPSEEK : llmType.getType());
    }

    @Override
    public ImageModel getImageModel() {
        return getImageModel(LlmType.OPENAI);
    }

    @Override
    public ImageModel getImageModel(LlmType llmType) {
        return imageModelMap.get(llmType.getType() == null ? LlmType.OPENAI : llmType.getType());
    }

    @Override
    public EmbeddingModel getEmbeddingModel() {
        return getEmbeddingModel(LlmType.OPENAI);
    }

    @Override
    public EmbeddingModel getEmbeddingModel(LlmType llmType) {
        return embeddingModelMap.get(llmType.getType() == null ? LlmType.OPENAI : llmType.getType());
    }

    @Override
    public void onApplicationEvent(ApplicationReadyEvent event) {
        Map<String, ChatModel> chatModelMapOfBeanMap = event.getApplicationContext().getBeansOfType(ChatModel.class);
        chatModelMap.putAll(chatModelMapOfBeanMap);

        Map<String, ImageModel> imageModelMapOfBeanMap = event.getApplicationContext().getBeansOfType(ImageModel.class);
        imageModelMap.putAll(imageModelMapOfBeanMap);

        Map<String, EmbeddingModel> embeddingModelMapOfBeanMap = event.getApplicationContext().getBeansOfType(EmbeddingModel.class);
        embeddingModelMap.putAll(embeddingModelMapOfBeanMap);
    }

}
